"""
@2024-12-12: 修复了make_stack方法中减扣本底时因数据类型导致的bug。大致情况是这样的，原先bias[y1:y2,:]的类型
            是无符号整形，将它作为“被减数”从32位浮点型变量img中减去时得到的结果出现了异常，处理后的电荷转移平场
            图像（原平场图像直接叠加，然后减去平均本底或者暗场）出现异常，overscan区域的数值应该基本为0，然而却
            出现了好几万的DN值，且有大量数值变化剧烈的点。
"""

import gc
import time
import psutil
from tqdm import trange
import numpy as np
from astropy.io import fits

import numpy as np
from .utils import get_block_limits
from .utils import timetag, hltxt, make_dir, clean_dir, remove_dir

def make_stack(flist, nclip=0, bias=None, norm=None,
               meanfunc=np.mean, dtype=np.float32, malloc=None):
    hltxt('**** calling make_stack, flist:\n{}\n'.format(flist))
    
    # memory management
    if malloc is None:
        malloc = psutil.virtual_memory().available / 1024 ** 3 / 3

    # open image files
    n = len(flist)
    fimgs = [fits.open(f) for f in flist]
    ny, nx = fimgs[0][0].data.shape

    # allocate data
    data = np.zeros((ny, nx), dtype=dtype)
    noise = np.zeros((ny, nx), dtype='float32')

    # estimate memory usage
    ncols = malloc / (nx * 8 / 1024 ** 3) / 2 / n
    if ncols < 2:
        raise Exception('Error: low memory.')
    nblocks = int(np.floor(ny / ncols)) + 1
    ylims = get_block_limits(ny, nblocks)

    # stack image
    if n == 1:
        data = fimgs[0][0].data.astype(dtype)
        if bias is not None:
            # hltxt('**** subtracting bias ****')
            data -= bias
        if norm is not None:
            # hltxt('**** multiplying norm ****')
            data *= norm[0]
    else:
        # read the image by blocks
        for i in trange(len(ylims) - 1, desc="processing {} images".format(n)):
            y1, y2 = ylims[i], ylims[i + 1]
            cutout = np.zeros((n, y2 - y1, nx), dtype=dtype)
            for j in range(n):
                img = fimgs[j][0].data[y1:y2, :]
                if bias is not None:
                    # img = img - bias[y1:y2, :] # 这一步的本底减扣出了问题
                    img = img - bias[y1:y2, :].astype(np.float32) # 这里做类型转换后就正常了！！！
                if norm is not None:
                    img = img * norm[j]
                cutout[j, :, :] = img
            noise[y1:y2, :] = np.std(cutout, axis=0, ddof=1)
            if nclip > 0:
                cutout.sort(axis=0)
                cutout = cutout[nclip:-nclip, :, :]
            data[y1:y2, :] = meanfunc(cutout, axis=0)
        del cutout, img

    for f in fimgs:
        f.close()
    gc.collect()

    return data, noise


def _get_nclip(n, clip_frac=0.05):
    if n <= 2:
        nclip = 0
    else:
        nclip = int(n * clip_frac) + 1
    return nclip


def stack_image(flist, outname, nclip=None, fbias=None,
                norm=None, normtag=None,
                dtype=None, header_info=None):
    if header_info is None:
        header_info = dict()
    t0 = time.time()
    hltxt(' **** stacking image {} ****\n'.format(outname))

    n = len(flist)
    if nclip is None:
        nclip = _get_nclip(len(flist))

    # bias
    bias = None
    if fbias is not None:
        hltxt('**** fbias is NOT none, loading fbias: {} ****'.format(fbias))
        bias = fits.getdata(fbias).astype(np.float32)   # 提前做好类型转换，防止后续处理步骤出现类型问题
        # bias = fits.getdata(fbias)

    # normalization
    if norm is None and normtag is not None:
        hltxt('**** getting norm')
        norm = np.asarray([1.0 / fits.getheader(f)[normtag] for f in flist])

    dtypeX = np.float32
    stacked, noise = make_stack(flist, nclip=nclip, bias=bias, norm=norm, dtype=dtypeX)

    hdulist = fits.HDUList([fits.PrimaryHDU(stacked), fits.ImageHDU(noise.data)])
    hdulist[1].name = 'NOISE'
    for hdu in hdulist:
        hdu.header['NSTACK'] = n - nclip * 2
        hdu.header['NCLIP'] = nclip
        for key in header_info:
            hdu.header[key] = header_info[key]
    print("stack output image",outname)
    hdulist.writeto(outname, overwrite=True)
    hltxt('image stacked {}, execution time {:.2f} s'.format(
        outname, time.time() - t0))

    return True


def stack_median(flist, outname, fbias=None, dtype=None, header_info=None):
    if header_info is None:
        header_info = dict()
    t0 = time.time()
    hltxt(' ***** stacking median {} ****\n'.format(outname))

    # bias
    if fbias is not None:
        bias = fits.getdata(fbias)
    else:
        bias = 0

    # dtype
    if dtype is None:
        # dtype = fits.getdata(flist[0]).dtype    
        dtype = np.float32     #revised @20241213

    stacked, _ = make_stack(flist, meanfunc=np.median, dtype=dtype)

    hdu = fits.PrimaryHDU(stacked - bias)
    for key in header_info:
        hdu.header[key] = header_info[key]
    hdu.writeto(outname, overwrite=True)
    hltxt('median image {} generated, execution time {:.2f} s'.format(
        outname, time.time() - t0))

    return True
